#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RRT改进模块

此模块提供RRT算法的改进实现，
包括RRT*、信息引导RRT*和路径平滑等增强功能。
"""

import numpy as np
import time
from scipy.spatial.transform import Rotation

from .rrt_base import RRTBase
from .node import Node


class RRTImproved(RRTBase):
    """
    带有RRT*和信息引导采样等优化的改进RRT实现。
    """
    
    def __init__(self, collision_checker, config):
        """
        使用碰撞检查器和配置初始化改进的RRT。
        
        参数:
            collision_checker: 碰撞检查器对象
            config (dict): 配置字典
        """
        super().__init__(collision_checker, config)
        
        # 提取RRT*特定的配置参数
        self.optimization_config = config.get("optimization", {})
        self.enable_smoothing = self.optimization_config.get("enable_smoothing", True)
        self.max_smooth_iterations = self.optimization_config.get("max_smooth_iterations", 100)
        self.rewire_radius = self.optimization_config.get("rewire_radius", 15.0)
        self.use_informed_sampling = self.optimization_config.get("use_informed_sampling", True)
        
        # 探索策略参数
        self.explore_strategy = config.get("explore_strategy", "random")
        self.adaptive_config = config.get("adaptive", {})
        self.adaptive_bias_min = self.adaptive_config.get("bias_min", 0.05)
        self.adaptive_bias_max = self.adaptive_config.get("bias_max", 0.8)
        self.adaptive_radius_factor = self.adaptive_config.get("radius_factor", 2.0)
        self.voronoi_bias = self.adaptive_config.get("voronoi_bias", 0.3)
        
        # 信息引导采样状态
        self.informed_sampling_active = False
        self.c_best = float('inf')  # 当前到目标的最佳成本
        self.ellipse_params = None  # 信息引导采样椭圆的参数
        
        # 初始化自适应采样参数
        self.current_goal_bias = self.goal_bias
        self.narrow_passage_regions = []
        self.collision_history = []
        self.exploration_phase = True  # 从探索阶段开始
        
    def plan(self, start, goal, max_iterations=None, goal_tolerance=None, plotter=None):
        """
        使用改进的RRT规划从起点到目标的路径。
        
        参数:
            start (tuple): 起始位置 (x, y)
            goal (tuple): 目标位置 (x, y)
            max_iterations (int, optional): 最大迭代次数
            goal_tolerance (float, optional): 目标达成容差
            plotter (Plotter, optional): 用于可视化的绘图器
            
        返回:
            tuple: (path, info) 其中path是点的列表，info是一个字典
        """
        # 运行带有改进的基础RRT规划算法
        path, info = super().plan(start, goal, max_iterations, goal_tolerance, plotter)
        
        # 如果我们找到了路径且启用了平滑，则平滑路径
        if path and len(path) > 2 and self.enable_smoothing:
            path = self._smooth_path(path)
            
            # 更新info中的路径
            info["smoothed_path"] = path
            
            # 如果提供了绘图器，则进行最终可视化更新
            if plotter:
                plotter.update_planning(self.tree, self.nearest_to_goal, self.goal_node, path)
                
        return path, info
    
    def _sample_free(self):
        """
        在自由空间中采样随机点，使用各种策略。
        
        返回:
            tuple: 采样点的(x, y)坐标
        """
        # 如果信息引导采样已激活且已启用，则使用它
        if self.use_informed_sampling and self.informed_sampling_active and self.ellipse_params:
            # 尝试信息引导采样
            point = self._informed_sample()
            
            # 验证点是否有效（在地图内且无碰撞）
            map_instance = self.collision_checker.map
            robot_radius = self.collision_checker.robot.radius
            
            if point is not None and map_instance.is_valid_point(point[0], point[1], robot_radius):
                return point
        
        # 如果配置了自适应采样策略，则应用它
        if self.explore_strategy == "adaptive":
            # 根据规划进度更新自适应采样参数
            self._update_adaptive_parameters()
            
            # 根据当前状态在不同的采样策略之间选择
            sampling_choice = np.random.random()
            
            if sampling_choice < self.current_goal_bias:
                # 带有当前自适应偏差的目标偏向采样
                return self._goal_biased_sample()
            elif sampling_choice < self.current_goal_bias + self.voronoi_bias and len(self.narrow_passage_regions) > 0:
                # Voronoi/窄通道偏向采样
                return self._voronoi_biased_sample()
            else:
                # 带有碰撞感知的均匀随机采样
                return self._collision_aware_sample()
        elif self.explore_strategy == "goal_biased":
            # 父类中的默认目标偏向采样
            return super()._sample_free()
        else:
            # 父类中的默认随机采样
            return super()._sample_free()
    
    def _update_adaptive_parameters(self):
        """
        根据规划进度更新自适应采样参数。
        """
        # 如果我们有足够的节点，开始识别窄通道
        if self.tree.size > 50 and len(self.collision_history) > 20:
            # 基于碰撞历史识别窄通道
            collision_rate = sum(self.collision_history[-20:]) / 20.0
            
            # 根据碰撞率调整目标偏差
            # 更高的碰撞率 → 更多探索（较低的目标偏差）
            # 较低的碰撞率 → 更多开发（较高的目标偏差）
            if collision_rate > 0.7:  # 高碰撞环境
                self.current_goal_bias = self.adaptive_bias_min
                self.exploration_phase = True
            elif collision_rate < 0.3:  # 低碰撞环境
                self.current_goal_bias = self.adaptive_bias_max
                self.exploration_phase = False
            else:
                # 在最小值和最大值之间线性插值
                self.current_goal_bias = self.adaptive_bias_min + \
                    (self.adaptive_bias_max - self.adaptive_bias_min) * (1 - collision_rate)
        
        # 如果我们有足够的信息，动态更新重连接半径
        if self.tree.size > 100:
            # 根据节点密度调整重连接半径
            # 我们希望在稀疏区域有更大的半径，在密集区域有更小的半径
            map_width, map_height = self.collision_checker.map.get_dimensions()
            map_area = map_width * map_height
            node_density = self.tree.size / map_area
            
            # 基于节点密度和树大小的重连接半径
            self.rewire_radius = max(
                5.0,  # 最小半径
                min(
                    30.0,  # 最大半径
                    self.adaptive_radius_factor * np.sqrt(np.log(self.tree.size) / node_density)
                )
            )
    
    def _goal_biased_sample(self):
        """
        生成偏向目标的样本。
        
        返回:
            tuple: 采样点的(x, y)坐标
        """
        # 获取目标位置
        goal_node = self.goal_node
        if goal_node is None and self.nearest_to_goal is not None:
            goal_pos = self.nearest_to_goal.position
        else:
            # 退回到随机采样
            return super()._sample_free()
        
        # 从以目标为中心的高斯分布采样
        std_dev = self.rewire_radius  # 使用重连接半径作为标准差
        
        # 采样直到获得有效点
        map_instance = self.collision_checker.map
        robot_radius = self.collision_checker.robot.radius
        
        for _ in range(10):  # 尝试最多10次
            x = np.random.normal(goal_pos[0], std_dev)
            y = np.random.normal(goal_pos[1], std_dev)
            
            # 确保点在地图边界内
            width, height = map_instance.get_dimensions()
            x = np.clip(x, 0, width)
            y = np.clip(y, 0, height)
            
            # 检查点是否无碰撞
            if map_instance.is_valid_point(x, y, robot_radius):
                return (x, y)
        
        # 如果所有尝试都失败，退回到常规采样
        return super()._sample_free()
    
    def _voronoi_biased_sample(self):
        """
        使用Voronoi偏差生成偏向窄通道的样本。
        
        返回:
            tuple: 采样点的(x, y)坐标
        """
        # 如果我们已识别窄通道，向它们偏向采样
        if len(self.narrow_passage_regions) > 0:
            # 随机选择一个窄通道区域
            region = self.narrow_passage_regions[np.random.randint(0, len(self.narrow_passage_regions))]
            
            # 从区域周围较小的分布中采样
            x = np.random.uniform(region[0] - region[2], region[0] + region[2])
            y = np.random.uniform(region[1] - region[2], region[1] + region[2])
            
            # 确保点在地图边界内
            map_instance = self.collision_checker.map
            width, height = map_instance.get_dimensions()
            x = np.clip(x, 0, width)
            y = np.clip(y, 0, height)
            
            return (x, y)
        
        # 如果未识别到窄通道，退回到常规采样
        return super()._sample_free()
    
    def _collision_aware_sample(self):
        """
        生成碰撞历史感知的样本。
        
        返回:
            tuple: 采样点的(x, y)坐标
        """
        # 获取基本随机样本
        point = super()._sample_free()
        
        # 记录此采样尝试的碰撞
        map_instance = self.collision_checker.map
        collision = map_instance.check_collision(point[0], point[1], self.collision_checker.robot.radius)
        self.collision_history.append(1 if collision else 0)
        
        # 限制历史大小
        if len(self.collision_history) > 100:
            self.collision_history = self.collision_history[-100:]
        
        # 如果我们遇到碰撞，检查这是否可能是窄通道
        if collision:
            # 寻找附近的有效点
            valid_points = []
            for i in range(8):
                # 在8个方向上检查
                angle = i * np.pi / 4
                dx = np.cos(angle) * self.collision_checker.robot.radius * 2
                dy = np.sin(angle) * self.collision_checker.robot.radius * 2
                
                test_point = (point[0] + dx, point[1] + dy)
                if not map_instance.check_collision(test_point[0], test_point[1], 
                                                 self.collision_checker.robot.radius):
                    valid_points.append(test_point)
            
            # 如果我们在附近找到有效点，但点本身在碰撞中，
            # 这可能是窄通道
            if len(valid_points) >= 2:
                # 将此存储为潜在的窄通道
                region = (point[0], point[1], self.collision_checker.robot.radius * 3)
                
                # 检查此区域是否与现有区域重叠
                is_new_region = True
                for existing_region in self.narrow_passage_regions:
                    dist = np.sqrt((existing_region[0] - region[0])**2 + 
                                  (existing_region[1] - region[1])**2)
                    if dist < existing_region[2] + region[2]:
                        is_new_region = False
                        break
                
                if is_new_region:
                    self.narrow_passage_regions.append(region)
                    
                    # 限制窄通道区域的数量
                    if len(self.narrow_passage_regions) > 10:
                        self.narrow_passage_regions.pop(0)
                
                # 返回其中一个有效点
                return valid_points[np.random.randint(0, len(valid_points))]
        
        # 如果没有碰撞或找不到更好的点，则返回原始点
        if not collision:
            return point
        
        # 如果我们在这里，我们遇到了碰撞但找不到更好的点
        # 再尝试几次
        for _ in range(5):
            new_point = super()._sample_free()
            if not map_instance.check_collision(new_point[0], new_point[1], 
                                             self.collision_checker.robot.radius):
                return new_point
        
        # 最后手段，返回我们得到的任何点
        return super()._sample_free()
    
    def _informed_sample(self):
        """
        从信息子集（椭圆）中采样点。
        
        返回:
            tuple: 采样点的(x, y)坐标
        """
        # 提取椭圆参数
        c_min = self.ellipse_params["c_min"]
        center = self.ellipse_params["center"]
        C = self.ellipse_params["C"]
        
        # 使用拒绝采样从单位球体采样
        unit_ball_point = None
        while unit_ball_point is None:
            # 从单位正方形采样
            x = np.random.uniform(-1, 1)
            y = np.random.uniform(-1, 1)
            
            # 检查点是否在单位球内
            if x**2 + y**2 <= 1:
                unit_ball_point = np.array([x, y, 0])  # 添加z=0用于2D
        
        # 将点变换到椭圆
        ellipse_point = center + C @ unit_ball_point
        
        return ellipse_point[0], ellipse_point[1]
    
    def _check_segment_collision(self, from_point, to_point, resolution=0.1):
        """
        检查两点之间的线段是否无碰撞。
        增强了更高的分辨率以提高准确性。
        
        参数:
            from_point (tuple): 起始点 (x, y)
            to_point (tuple): 终点 (x, y)
            resolution (float): 沿线段进行碰撞检查的分辨率
            
        返回:
            bool: 若有碰撞则为True，若无碰撞则为False
        """
        # 使用稍高的分辨率以获得更好的准确性
        return super()._check_segment_collision(from_point, to_point, resolution * 0.75)
    
    def _setup_informed_sampling(self, start, goal, c_best):
        """
        设置信息引导采样参数。
        
        参数:
            start (tuple): 起始位置 (x, y)
            goal (tuple): 目标位置 (x, y)
            c_best (float): 当前最佳成本（路径长度）
        """
        self.informed_sampling_active = True
        self.c_best = c_best
        
        # 计算椭圆参数
        start_point = np.array([start[0], start[1], 0])  # 添加z=0用于2D
        goal_point = np.array([goal[0], goal[1], 0])     # 添加z=0用于2D
        
        # 计算最小距离（直线）
        c_min = np.linalg.norm(start_point - goal_point)
        
        # 计算椭圆中心
        center = (start_point + goal_point) / 2.0
        
        # 计算椭圆旋转和缩放
        a = c_best / 2.0  # 半长轴
        if c_min < 1e-6:  # 避免除以零
            c_min = 1e-6
        b = np.sqrt(a**2 - (c_min/2.0)**2)  # 半短轴
        
        # 椭圆变换矩阵
        if c_min < 1e-6:
            # 如果起点和目标非常接近，使用单位变换
            C = np.eye(3)
        else:
            # 计算旋转矩阵
            dir_vec = (goal_point - start_point) / c_min
            # 对于2D，我们只需要绕z轴旋转
            z_axis = np.array([0, 0, 1])
            # 找到从[1,0,0]到dir_vec在X-Y平面的旋转
            x_axis = np.array([1, 0, 0])
            # 将dir_vec投影到X-Y平面
            dir_vec_xy = np.array([dir_vec[0], dir_vec[1], 0])
            dir_vec_xy = dir_vec_xy / np.linalg.norm(dir_vec_xy)
            # 计算x_axis和dir_vec_xy之间的角度
            cos_angle = np.dot(x_axis, dir_vec_xy)
            angle = np.arccos(np.clip(cos_angle, -1.0, 1.0))
            if dir_vec_xy[1] < 0:  # 检查符号以确定正确的旋转方向
                angle = -angle
            
            # 创建旋转矩阵
            rotation = Rotation.from_rotvec(angle * z_axis)
            rot_matrix = rotation.as_matrix()
            
            # 创建缩放矩阵
            scale_matrix = np.diag([a, b, 1])
            
            # 组合旋转和缩放
            C = rot_matrix @ scale_matrix
        
        # 存储椭圆参数
        self.ellipse_params = {
            "c_min": c_min,
            "c_best": c_best,
            "center": center,
            "C": C
        }
    
    def _update_informed_sampling(self, c_best):
        """
        当找到更好的路径时更新信息引导采样参数。
        
        参数:
            c_best (float): 新的最佳成本（路径长度）
        """
        if self.ellipse_params:
            # 更新c_best和相关参数
            self.c_best = c_best
            self.ellipse_params["c_best"] = c_best
            
            # 更新半长轴
            a = c_best / 2.0
            c_min = self.ellipse_params["c_min"]
            
            # 重新计算半短轴
            if c_min < 1e-6:  # 避免除以零
                c_min = 1e-6
            b = np.sqrt(a**2 - (c_min/2.0)**2)
            
            # 更新变换矩阵中的缩放
            C = self.ellipse_params["C"]
            # 从当前C中提取旋转
            U, _, Vt = np.linalg.svd(C)
            rot_matrix = U @ Vt
            
            # 创建新的缩放矩阵
            scale_matrix = np.diag([a, b, 1])
            
            # 更新C
            C = rot_matrix @ scale_matrix
            self.ellipse_params["C"] = C
    
    def _smooth_path(self, path):
        """
        使用捷径技术平滑路径。
        
        参数:
            path (list): (x, y)坐标列表
            
        返回:
            list: 平滑后的路径
        """
        if len(path) <= 2:
            return path
        
        smoothed_path = path.copy()
        
        # 执行捷径迭代
        for _ in range(self.max_smooth_iterations):
            # 如果路径只有2个点，则无法进一步平滑
            if len(smoothed_path) <= 2:
                break
                
            # 选择两个随机索引
            idx1 = np.random.randint(0, len(smoothed_path) - 1)
            idx2 = np.random.randint(idx1 + 1, len(smoothed_path))
            
            # 检查是否可以创建捷径
            if idx2 - idx1 <= 1:
                continue  # 相邻点，无需平滑
                
            # 尝试直接连接两点
            point1 = smoothed_path[idx1]
            point2 = smoothed_path[idx2]
            
            if not self._check_segment_collision(point1, point2):
                # 如果无碰撞，移除中间点
                smoothed_path = smoothed_path[:idx1 + 1] + smoothed_path[idx2:]
        
        return smoothed_path
    
    def _find_near_nodes(self, node):
        """
        查找给定节点重连接半径内的节点。
        
        参数:
            node (Node): 要查找邻居的节点
            
        返回:
            list: 附近节点的列表
        """
        return self.tree.get_nodes_in_radius(node.x, node.y, self.rewire_radius)
    
    def _compute_new_cost(self, from_node, to_node):
        """
        计算通过'from_node'到达'to_node'的成本。
        
        参数:
            from_node (Node): 源节点
            to_node (Node): 目标节点
            
        返回:
            float: 新成本
        """
        # 成本 = 父节点的成本 + 到父节点的距离
        distance = self._compute_distance(from_node.position, to_node.position)
        return from_node.cost + distance 